#include <kalman_node.h>

using namespace KalmanJoint;
using namespace std;

double radians(double degrees) {
    return degrees * M_PI / 180.0;
}

double wrap_to_pi(double angle){
    angle = fmod(angle + M_PI, 2 * M_PI);
    if (angle < 0){angle += 2 * M_PI;}
    return angle - M_PI;
}

Eigen::Matrix<double, 3, 1> calc_oil_cylinder_length(double boom_to_body, double forearm_to_boom, double bucket_to_forearm){
    Eigen::Matrix<double, 3, 1> cylinder_status;

    double angle_ACB = radians(deg_TCA) + radians(deg_BCF) + radians(boom_to_body);
    angle_ACB = wrap_to_pi(angle_ACB);
    cylinder_status(0, 0) = sqrt(pow(mm_AC, 2) + pow(mm_BC, 2) - 2*mm_AC*mm_BC*cos(angle_ACB));

    double phi = M_PI - radians(deg_DFC) - radians(deg_QFG) - radians(deg_GFE) - radians(forearm_to_boom);
    phi = wrap_to_pi(phi);
    cylinder_status(1, 0) = sqrt(pow(mm_DF, 2) + pow(mm_EF, 2) - 2*mm_EF*mm_DF*cos(phi));

    double angle_KQN = M_PI - radians(deg_NQF) - radians(bucket_to_forearm) - radians(deg_KQV);
    angle_KQN = wrap_to_pi(angle_KQN);
    double distance_NK = sqrt(pow(mm_QN, 2) + pow(mm_QK, 2)  - 2*mm_QK*mm_QN*cos(angle_KQN));
    double angle_MNK = acos((pow(mm_MN, 2) + pow(distance_NK,2) - pow(mm_MK, 2))/(2*mm_MN*distance_NK));
    double angle_QNK = asin(mm_QK*sin(angle_KQN)/distance_NK);
    double angle_GNM = M_2_PI - radians(deg_GNF) - radians(deg_FNQ) - angle_MNK - angle_QNK;
    angle_GNM = wrap_to_pi(angle_GNM);
    cylinder_status(2, 0) = sqrt(pow(mm_GN, 2) + pow(mm_MN, 2) - 2*mm_MN*mm_GN*cos(angle_GNM));

    return cylinder_status;
}

Eigen::Matrix<double, 3, 1> calc_bucket_position(double boom,double forearm,double bucket,double slewing){
    Eigen::Matrix<double, 3, 1> bucket_position;

    boom = radians(boom);
    forearm = radians(forearm);
    bucket = radians(bucket);
    slewing = radians(slewing);

    bucket_position(0, 0) = cos(slewing)*(mm_QV*cos(bucket) + mm_FQ*cos(forearm) + mm_CF*cos(boom) + mm_JC);
    bucket_position(1, 0) = sin(slewing)*(mm_QV*cos(bucket) + mm_FQ*cos(forearm) + mm_CF*cos(boom) + mm_JC);
    bucket_position(2, 0) = mm_QV*sin(bucket) +mm_FQ*sin(forearm) + mm_CF*sin(boom) + mm_JL;

    return bucket_position;
}

// RGBD数据回调函数
void rgbd_apriltag_callback(const cv_joint_angle::Forearm::ConstPtr& msg, MultiSensorKalmanFilter& joint_kalman){
    if(joint_kalman.getFilterInit()){
        Eigen::VectorXd measurement(3);
        measurement << 0, msg->forearm, 0;
        long double timestamp = msg->header.stamp.sec * 1000. + msg->header.stamp.nsec / 1000000.;
        if(timestamp <= joint_kalman.getNewestTimeStamp(RGBD_Apriltag)){
            ROS_WARN_THROTTLE(10, "The forearm RGBD timestamp is late");
            return;
        }
		if(fabs(msg->forearm-joint_kalman.getStateEstimate()(1, 0))>20){
            ROS_WARN_THROTTLE(10, "The gap between forearm RGBD and filter value is large");
            joint_kalman.setUsableFlag(RGBD_Apriltag, false);
        }
        joint_kalman.setUsableFlag(RGBD_Apriltag, true);
        joint_kalman.setNewestData(RGBD_Apriltag, measurement);
        joint_kalman.setNewestTimeStamp(RGBD_Apriltag, timestamp);
    }
}

void rgbd_RFRNN_callback(const cv_joint_angle::Bucket::ConstPtr& msg, MultiSensorKalmanFilter& joint_kalman){
    if(joint_kalman.getFilterInit()){
        Eigen::VectorXd measurement(3);
        measurement << 0, 0, msg->bucket;
        long double timestamp = msg->header.stamp.sec * 1000. + msg->header.stamp.nsec / 1000000.;
        if(timestamp <= joint_kalman.getNewestTimeStamp(RGBD_YOLO_RFRNN)){
            ROS_WARN_THROTTLE(10, "The RFRNN timestamp is late");
            return;
        }
		if(fabs(msg->bucket-joint_kalman.getStateEstimate()(2, 0))>20){
            ROS_WARN_THROTTLE(10, "The gap between bucket RGBD and filter value is large");
            joint_kalman.setUsableFlag(RGBD_YOLO_RFRNN, false);
        }
        joint_kalman.setUsableFlag(RGBD_YOLO_RFRNN, true);
        joint_kalman.setNewestData(RGBD_YOLO_RFRNN, measurement);
        joint_kalman.setNewestTimeStamp(RGBD_YOLO_RFRNN, timestamp);
    }
}

// IMU数据回调函数、滤波器更新与消息发布
void inclination_angle_callback(const mavros_msgs::ExcavatorInclination::ConstPtr& msg, MultiSensorKalmanFilter& joint_kalman,
        ros::Publisher& joint_kalman_pub){
    ROS_INFO_ONCE("kalman node get inclination info!");
    Eigen::Vector3d filter_angle;
    kalman_joint::KalmanJoint filter_msg;
    static long double last_time_ms = msg->header.stamp.sec * 1000. + msg->header.stamp.nsec / 1000000.;
    static long double last_error_time_ms = msg->header.stamp.sec * 1000. + msg->header.stamp.nsec / 1000000.;
    const long double curr_time_ms = msg->header.stamp.sec * 1000. + msg->header.stamp.nsec / 1000000.;

    // 更新数据
    Eigen::VectorXd curr_angle(3, 1), curr_vel(3, 1);
    curr_angle << msg->boom_deg, msg->forearm_deg, msg->bucket_deg;
    curr_vel  << msg->boom_vel, msg->forearm_vel, msg->bucket_vel;

    auto publishError = [&](const std::string& status, const std::string& warning_msg) {
        joint_kalman.setFilterInit(false);
        filter_msg.header = msg->header;
        filter_msg.status = status;
        filter_msg.angle_boom = curr_angle(0, 0);
        filter_msg.angle_forearm = curr_angle(1, 0);
        filter_msg.angle_bucket = curr_angle(2, 0);
        joint_kalman_pub.publish(filter_msg);
        ROS_WARN_STREAM_DELAYED_THROTTLE(2, warning_msg);
        last_error_time_ms = curr_time_ms;
    };

    if (curr_time_ms - last_error_time_ms > 500) {
        publishError("bad_stamp", "message stamp incorrect: old time: "
                     + std::to_string(last_error_time_ms) + ", new time: " + std::to_string(curr_time_ms) + "!!!");
        return;
    }

    for (int i = 0; i < 2; i++) {
        if (curr_angle(i, 0) > 180 || curr_angle(i, 0) < -180) {
            publishError("bad_msg", "message info incorrect!!!");
            return;
        }
    }
    joint_kalman.updateAllData(IMU, curr_angle, curr_time_ms);

   auto checkSensorTimestamp = [&](SensorNumber sensor, const std::string& sensor_name) {
        if(joint_kalman.getUsableFlag(sensor)){
            if (fabs(curr_time_ms - joint_kalman.getSensorTimeStamp(sensor)) > 3000) {
                if (joint_kalman.getSensorTimeStamp(sensor) != 0.) {
                    ROS_WARN_STREAM_DELAYED_THROTTLE(2, sensor_name + " message stamp incorrect: old time: "
                                                    + std::to_string(joint_kalman.getSensorTimeStamp(sensor)) +
                                                    ", new time: " + std::to_string(curr_time_ms) + "!!!");
                }
                joint_kalman.setUsableFlag(sensor, false);
            }
        }
    };

    checkSensorTimestamp(RGBD_YOLO_RFRNN, "RF_RNN");
    checkSensorTimestamp(RGBD_Apriltag, "Apriltag");

    // 根据可用数据滤波
    filter_msg.header = msg->header;
    if(joint_kalman.getFilterInit()){
        // 目前异步滤波一次callback仅融合两个数据源, 优先级:rfrnn->apriltag->inclination
        if (joint_kalman.getUpdateFlag(RGBD_YOLO_RFRNN) && joint_kalman.getUsableFlag(RGBD_YOLO_RFRNN)) {
            ROS_INFO_ONCE("kalman node get RFRNN info!");
            joint_kalman.updateFromHis(RGBD_YOLO_RFRNN);
            filter_msg.status = "RFRNN flitered";
        }
        else if (joint_kalman.getUpdateFlag(RGBD_Apriltag) && joint_kalman.getUsableFlag(RGBD_Apriltag)) {
            ROS_INFO_ONCE("kalman node get apriltag info!");
            joint_kalman.updateFromHis(RGBD_Apriltag);
            filter_msg.status = "RGBD forearm flitered";
        }
        joint_kalman.setInput(curr_vel);
        joint_kalman.setTimeStep(last_time_ms, curr_time_ms);
        joint_kalman.setSensorData(IMU, curr_angle);
        joint_kalman.updateAllSensor();
        if (!(joint_kalman.getUpdateFlag(RGBD_YOLO_RFRNN) || joint_kalman.getUpdateFlag(RGBD_Apriltag))){
            filter_msg.status = "IMU flitered";
        }
        joint_kalman.setUpdateFlag(RGBD_YOLO_RFRNN, false);
        joint_kalman.setUpdateFlag(RGBD_Apriltag, false);
        ROS_INFO_THROTTLE(300, "kalman running...");
    }else{
        joint_kalman.setFilterInit(true, curr_angle, IMU);
        filter_msg.status = "inited";
        ROS_INFO("kalman inited!");
    }

    filter_angle = joint_kalman.getStateEstimate();
    last_time_ms = curr_time_ms;
    last_error_time_ms = curr_time_ms;

    // 保存数据到队列中
    joint_kalman.addHisData(curr_time_ms);

    filter_msg.angle_boom = filter_angle(0,0);
    filter_msg.angle_forearm = filter_angle(1,0);
    filter_msg.angle_bucket = filter_angle(2,0);

    Eigen::Matrix<double, 3, 1> cylinder_length = calc_oil_cylinder_length(filter_angle(0,0), filter_angle(1,0) -filter_angle(0,0), filter_angle(2,0)-filter_angle(1,0));
    filter_msg.cylinder_boom = cylinder_length(0,0);
    filter_msg.cylinder_forearm = cylinder_length(1,0);
    filter_msg.cylinder_bucket = cylinder_length(2,0);

    Eigen::Matrix<double, 3, 1> bucket_position = calc_bucket_position(filter_angle(0,0), filter_angle(1,0), filter_angle(2,0), 0.);    //缺少回转数据
    filter_msg.bucket_x = bucket_position(0,0);
    filter_msg.bucket_y = bucket_position(1,0);
    filter_msg.bucket_z = bucket_position(2,0);

    Eigen::Matrix<double, 3, 3> cov = joint_kalman.getNoiseCov();
    filter_msg.cov_boom = cov(0,0);
    filter_msg.cov_forearm = cov(1,1);
    filter_msg.cov_bucket = cov(2,2);

    joint_kalman_pub.publish(filter_msg);
}

MultiSensorKalmanFilter kalman_init(){
    //创建多传感器参数结构体向量
    Eigen::MatrixXd imu_R(3,3);
    imu_R << 1., 0, 0,
                        0, 0.5, 0,
                        0, 0, 0.5;
    Eigen::MatrixXd imu_H(3,3);
    imu_H << 1., 0, 0,
                        0, 1., 0,
                        0, 0, 1.;     
    auto imu_params = MultiSensorKalmanFilter::CreateSensorParams(
                                                    SensorNumber::IMU, imu_H, imu_R);

    Eigen::MatrixXd apriltag_R(3,3);
    apriltag_R <<  0., 0, 0,
                                0, 1., 0,
                                0, 0,  0;     
    Eigen::MatrixXd apriltag_H(3,3);
    apriltag_H << 0, 0, 0,
                                0, 1., 0,
                                0, 0, 0;
    auto apriltag_params = MultiSensorKalmanFilter::CreateSensorParams(
                                                        SensorNumber::RGBD_Apriltag, apriltag_H, apriltag_R);   

    Eigen::MatrixXd rfrnn_R(3,3);
    rfrnn_R <<  0, 0, 0,
                                0, 0, 0,
                                0, 0, 1.;     
    Eigen::MatrixXd rfrnn_H(3,3);
    rfrnn_H << 0, 0, 0,
                                0, 0, 0,
                                0, 0, 1.;
    auto rfrnn_params = MultiSensorKalmanFilter::CreateSensorParams(
                                                        SensorNumber::RGBD_YOLO_RFRNN, rfrnn_H, rfrnn_R);   

    vector<MultiSensorKalmanFilter::SensorParams> sensors_params = {imu_params, apriltag_params, rfrnn_params};

    MultiSensorKalmanFilter joint_kalman = MultiSensorKalmanFilter(3, 3, sensors_params);

    Eigen::MatrixXd kf_A = Eigen::MatrixXd::Identity(3,3);
    joint_kalman.setTransitionMatrix(kf_A);

    Eigen::MatrixXd kf_B(3,3);
    kf_B <<   joint_kalman.getTimeStep(), 0, 0,
                        0, joint_kalman.getTimeStep(), 0,
                        0, 0, joint_kalman.getTimeStep();
    joint_kalman.setInputMatrix(kf_B);

    Eigen::MatrixXd kf_Q(3,3);
    kf_Q <<  1, 0, 0,
                        0, 0.5, 0,
                        0, 0, 0.5;
    joint_kalman.setProcessNoiseCov(kf_Q);

    // joint_kalman.setAlpha(0.3);

    return joint_kalman;
}

int main(int argc, char *argv[]){
    MultiSensorKalmanFilter joint_kalman = kalman_init();

    //创建ROS话题
    ros::init(argc, argv, "kalman_joint_node");
    ros::NodeHandle nh;

    // Retrieve parameters from the parameter server
    std::string pub_topic;
    std::string forearm_topic;
    std::string bucket_topic;
    std::string inclination_topic;

    nh.param<std::string>("pub_topic", pub_topic, "kalman_joint");
    nh.param<std::string>("forearm_topic", forearm_topic, "/cv_joint_angle/forearm_apriltag");
    nh.param<std::string>("bucket_topic", bucket_topic, "/cv_joint_angle/bucket_predict");
    nh.param<std::string>("inclination_topic", inclination_topic, "/mavros/excavator_inclination");

    // Set up publishers and subscribers
    ros::Publisher joint_kalman_pub = nh.advertise<kalman_joint::KalmanJoint>(pub_topic, 10);
    ros::Subscriber rgbd_forearm_sub = nh.subscribe<cv_joint_angle::Forearm>(
        forearm_topic, 10, boost::bind(rgbd_apriltag_callback, _1, boost::ref(joint_kalman)));
    ros::Subscriber rgbd_bucket_sub = nh.subscribe<cv_joint_angle::Bucket>(
        bucket_topic, 2, boost::bind(rgbd_RFRNN_callback, _1, boost::ref(joint_kalman)));
    ros::Subscriber inclination_angle_sub = nh.subscribe<mavros_msgs::ExcavatorInclination>(
        inclination_topic, 50, boost::bind(inclination_angle_callback, _1, boost::ref(joint_kalman), boost::ref(joint_kalman_pub)));

    ROS_INFO_ONCE("Kalman node started!");
    ros::spin();

    return 0;
}

